import pandas
import plotly.graph_objs as go
import plotly.offline as py
from collections import Counter
from combo.models.classifier_stacking import Stacking
from combo.models.classifier_comb import SimpleClassifierAggregator
from imblearn.over_sampling import BorderlineSMOTE
from scipy import stats
from sklearn.metrics import auc, roc_curve
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.svm import SVC
from sklearn.ensemble import GradientBoostingClassifier, RandomForestClassifier
from sklearn.neural_network import MLPClassifier
from sklearn.tree import DecisionTreeClassifier

from factors import (
    accounting_factors,
    market_factors,
    academic_factors,
)

factors = accounting_factors + market_factors + academic_factors


# Stage 1
stage1_models = {}
stage1_auc = {}
for i in "abcde":
    predicted = f"preFD_{i}"

    dataset = pandas.read_csv("./data/data.csv", sep="\t")
    # dataset = dataset[dataset['accper'] > "2017-01-01"]

    X = stats.zscore(dataset[factors])
    y = dataset[predicted]
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.3, random_state=1
    )

    oversample = BorderlineSMOTE()
    print(Counter(y_train))
    X_train, y_train = oversample.fit_resample(X_train, y_train)
    print(Counter(y_train))

    classifiers = [
        DecisionTreeClassifier(),
        LogisticRegression(penalty="none", solver="newton-cg"),
        MLPClassifier(hidden_layer_sizes=(100, 50, 10), max_iter=100000),
        SVC(probability=True),
        RandomForestClassifier(),
        GradientBoostingClassifier(),
    ]

    stage1_models[i] = SimpleClassifierAggregator(
        base_estimators=classifiers, method="average"
    )
    stage1_models[i].fit(X_train, y_train)
    prob = stage1_models[i].predict_proba(X_test)
    fpr, tpr, thresholds = roc_curve(y_test, prob[:, 1], pos_label=1)
    cur_auc = auc(fpr, tpr)
    stage1_auc[i] = cur_auc
    print(cur_auc)


# Stage 2
X = stats.zscore(dataset[["Size", "Lev", "PPE"]])
y = (
    dataset["preFD_a"]
    + dataset["preFD_b"]
    + dataset["preFD_c"]
    + dataset["preFD_d"]
    + dataset["preFD_e"]
)
y = y > 0
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=1)

oversample = BorderlineSMOTE()
print(Counter(y_train))
before = len(y_train)
print(before)
X_train, y_train = oversample.fit_resample(X_train, y_train)
print(Counter(y_train))
after = len(y_train)
print(after)

column_added = after - before
X_train = X_train.append(X_train[before:after])
y_train = y_train.append(y_train[before:after])
print(Counter(y_train))
X_train, y_train = oversample.fit_resample(X_train, y_train)
print(Counter(y_train))

X_train = X_train[after:]
y_train = y_train[after:]
